package com.my_test.webmagic.webmagic_modules;

import com.alibaba.fastjson.JSONObject;
import com.my_test.webmagic.module.entity.*;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkRequestBuilder;
import org.elasticsearch.client.transport.TransportClient;
import org.springframework.stereotype.Component;
import us.codecraft.webmagic.ResultItems;
import us.codecraft.webmagic.Task;
import us.codecraft.webmagic.pipeline.Pipeline;

import javax.annotation.Resource;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author zhangyalin@act-telecom.com
 * @since 2019-05-27 16:58
 **/
@Component
public class WebMagicEsPipeline implements Pipeline {

    @Resource
    private TransportClient transportClient;

    @Override
    public void process(ResultItems resultItems, Task task) {

        String url = resultItems.getRequest().getUrl();
        Object list0 = resultItems.get("entityList");

        BulkRequestBuilder bulk = transportClient.prepareBulk();
        Map<String, Object> map0 = new HashMap<>();
        if (url.matches(WebMagicProcessor.typeRegex)) {
            //类型
            List<TiebaType> list = (List<TiebaType>) list0;
            int m = 0;
            for (int i = 0; i < list.size(); i++) {
                m++;
                TiebaType entity = list.get(i);
                map0 = (Map<String, Object>) JSONObject.toJSON(entity);
                bulk.add(transportClient.prepareIndex("tiebadata", "type", entity.getId() + "").setSource(map0));
                if (i % 10 == 0) {
                    m = 0;
                    bulk.execute().actionGet();
                    bulk = transportClient.prepareBulk();
                }
            }
            if (m != 0) {
                bulk.execute().actionGet();
            }

        } else if (url.matches(WebMagicProcessor.tiebaRegex)) {
            //帖子
            List<Tiezi> list = (List<Tiezi>) list0;
            int m = 0;
            for (int i = 0; i < list.size(); i++) {
                m++;
                Tiezi entity = list.get(i);
                map0 = (Map<String, Object>) JSONObject.toJSON(entity);
                bulk.add(transportClient.prepareIndex("tiebadata", "tiezi", entity.getId() + "").setSource(map0));
                if (i % 10 == 0) {
                    m = 0;
                    bulk.execute().actionGet();
                    bulk = transportClient.prepareBulk();
                }
            }
            if (m != 0) {
                bulk.execute().actionGet();
            }
        } else if (url.matches(WebMagicProcessor.tieziRegex)) {
            //楼层
            List<Floor> list = (List<Floor>) list0;
            Floor entity;
            int m = 0;
            for (int i = 0; i < list.size(); i++) {
                m++;
                entity = list.get(i);
                map0 = (Map<String, Object>) JSONObject.toJSON(entity);
                bulk.add(transportClient.prepareIndex("tiebadata", "floor", entity.getId() + "").setSource(map0));
                if (i % 10 == 0) {
                    m = 0;
                    bulk.execute().actionGet();
                    bulk = transportClient.prepareBulk();
                }
            }
            if (m != 0) {
                bulk.execute().actionGet();
            }

            List<User> userList = resultItems.get("userList");
            BulkRequestBuilder userBulk = transportClient.prepareBulk();
            User user;
            Map<String, Object> map1;
            int n = 0;
            for (int i = 0; i < userList.size(); i++) {
                n++;
                user = userList.get(i);
                map1 = (Map<String, Object>) JSONObject.toJSON(user);
                userBulk.add(transportClient.prepareIndex("tiebadata", "user", user.getId() + "").setSource(map1));
                if (i % 10 == 0) {
                    n = 0;
                    userBulk.execute().actionGet();
                    userBulk = transportClient.prepareBulk();
                }
            }
            if (n != 0) {
                BulkRequest request = userBulk.request();
                request.requests().size();
                userBulk.execute().actionGet();
            }

        } else if (url.matches(WebMagicProcessor.type2Regex)) {
            //贴吧
            List<Tieba> list = (List<Tieba>) list0;
            int m = 0;
            for (int i = 0; i < list.size(); i++) {
                m++;
                Tieba entity = list.get(i);
                map0 = (Map<String, Object>) JSONObject.toJSON(entity);
                bulk.add(transportClient.prepareIndex("tiebadata", "tieba", entity.getId() + "").setSource(map0));
                if (i % 10 == 0) {
                    m = 0;
                    bulk.execute().actionGet();
                    bulk = transportClient.prepareBulk();
                }
            }
            if (m != 0) {
                bulk.execute().actionGet();
            }
        }
    }
}
